import numpy as np
import heapq

class Node:
    def __init__(self, x, y, cost, parent_idx):
        self.x = x  # x coordinate
        self.y = y  # y coordinate
        self.cost = cost  # total cost to reach this node
        self.parent_idx = parent_idx  # index of parent node

    def __lt__(self, other):
        return self.cost < other.cost

class AStarPlanner:
    def __init__(self, grid_map, resolution=1.0, robot_radius=1.0):
        self.grid_map = grid_map  # 100x100 grid map
        self.resolution = resolution
        self.robot_radius = robot_radius
        self.motion = self.get_motion_model()
        self.cost_map = np.zeros_like(grid_map)

    def planning(self, start, goal):
        """
        A* path planning algorithm
        :param start: start position [x, y]
        :param goal: goal position [x, y]
        :return: path coordinates [[x1,y1], [x2,y2], ...]
        """
        # Convert world coordinates to grid indices
        start_idx = self.world2grid(start)
        goal_idx = self.world2grid(goal)

        # Initialize open and closed sets
        open_set = []
        closed_set = dict()

        # Start node
        start_node = Node(start_idx[0], start_idx[1], 0.0, -1)
        heapq.heappush(open_set, (start_node.cost, start_node.x, start_node.y, id(start_node)))
        closed_set[(start_node.x, start_node.y)] = start_node

        # Main loop
        while open_set:
            cost, x, y, node_id = heapq.heappop(open_set)
            current = closed_set[(x, y)]

            # Check if goal is reached
            if (x, y) == (goal_idx[0], goal_idx[1]):
                return self.extract_path(closed_set, current)

            # Explore neighbors
            for i, _ in enumerate(self.motion):
                dx, dy, move_cost = self.motion[i]
                nx = x + dx
                ny = y + dy

                # Check if node is within grid bounds
                if not self.is_in_grid(nx, ny):
                    continue

                # Check if node is obstacle-free
                if self.grid_map[ny, nx] == 1:
                    continue

                # Calculate new cost
                new_cost = current.cost + move_cost
                new_node = Node(nx, ny, new_cost, (x, y))

                # Add to open set if not in closed set or has lower cost
                if (nx, ny) not in closed_set or new_cost < closed_set[(nx, ny)].cost:
                    heapq.heappush(open_set, (new_node.cost + self.heuristic(nx, ny, goal_idx), nx, ny, id(new_node)))
                    closed_set[(nx, ny)] = new_node

        return None  # No path found

    def heuristic(self, x, y, goal):
        """Heuristic function using Euclidean distance"""
        return np.hypot(x - goal[0], y - goal[1])

    def is_in_grid(self, x, y):
        """Check if coordinates are within grid bounds"""
        return 0 <= x < self.grid_map.shape[1] and 0 <= y < self.grid_map.shape[0]

    def world2grid(self, world_coords):
        """Convert world coordinates to grid indices"""
        return [int(world_coords[0] / self.resolution), int(world_coords[1] / self.resolution)]

    def grid2world(self, grid_indices):
        """Convert grid indices to world coordinates"""
        return [grid_indices[0] * self.resolution, grid_indices[1] * self.resolution]

    def extract_path(self, closed_set, goal_node):
        """Extract path from closed set"""
        path = []
        current = goal_node
        while current.parent_idx != -1:
            path.append(self.grid2world([current.x, current.y]))
            current = closed_set[current.parent_idx]
        path.append(self.grid2world([current.x, current.y]))
        return path[::-1]  # Reverse to get path from start to goal

    @staticmethod
    def get_motion_model():
        """Define possible motion directions with costs"""
        # dx, dy, cost
        motion = [
            [1, 0, 1.0],   # Right
            [0, 1, 1.0],   # Up
            [-1, 0, 1.0],  # Left
            [0, -1, 1.0],  # Down
            [1, 1, 1.414], # Right-Up
            [1, -1, 1.414],# Right-Down
            [-1, 1, 1.414],# Left-Up
            [-1, -1, 1.414] # Left-Down
        ]
        return motion